# -*- coding: utf-8 -*-
# @Author: bao
# @Date:   2017-03-31 15:10:31
# @Last Modified by:   bao
# @Last Modified time: 2017-04-06 16:31:30
import sys
sys.path.append("..")

import re
import random
from CoreSim.MScheduler import *
from CoreSim.Flow import *
from CoreSim.Event import *
from CoreSim.Link import *
from CoreSim.Unit import *

def assignNewFlow(srcHost, dstHost, flows, flowSize, startTime, mcastFlowId, numHosts, numToRs, numCores):
    # unicast flow
    flow = Flow()
    flow.startId = srcHost
    flow.endId = dstHost
    flow.SetFlowSize(flowSize)
    flow.startTime = startTime
    flow.flowId = len(flows) + 1
    flow.mcastFlowId = mcastFlowId
    flows.append(flow)
    
    # assign path
    srcToRId = GetToRId(srcHost, numHosts)
    dstToRId = GetToRId(dstHost, numHosts)
    srcToRNodeId = ConvertToNodeId(srcToRId, ToR, numHosts*numToRs, numToRs)
    dstToRNodeId = ConvertToNodeId(dstToRId, ToR, numHosts*numToRs, numToRs)
    
    if srcToRNodeId == dstToRNodeId:
        pathNodeIds = [flow.startId, srcToRNodeId, flow.endId]
    else:
        coreId = random.randint(1, numCores)
        coreNodeId = ConvertToNodeId(coreId, Core, numHosts*numToRs, numToRs)
        pathNodeIds = [flow.startId, srcToRNodeId, coreNodeId ,dstToRNodeId, flow.endId]
    
    flow.BuildPath(pathNodeIds)

    return flow

class MSchedulerDoubling(MScheduler):
    '''
    Multicast scheduler
    '''
    def __init__(self, topo_args, flow_args):
        MScheduler.__init__(self, topo_args, flow_args)
        self.mcastSubtree = {} # binomial tree
    
    def InitialMulticast(self, eventQueue):
        for mflowId in self.mcastFlows:
            mcastFlow = self.mcastFlows[mflowId]
            srcHost = mcastFlow.srcHost
            index = len(mcastFlow.dstHosts.keys())/2
            dstHost = mcastFlow.dstHosts.keys()[index]

            self.mcastSubtree[mflowId] = {}
            self.mcastSubtree[mflowId][srcHost] = mcastFlow.dstHosts.keys()[:index]
            self.mcastSubtree[mflowId][dstHost] = mcastFlow.dstHosts.keys()[index+1:]
            print srcHost, dstHost
            print mcastFlow.dstHosts.keys()
            print 'mcastSubtree', self.mcastSubtree
            
            flow = assignNewFlow(srcHost, dstHost, self.flows, mcastFlow.flowSize, mcastFlow.startTime, \
                mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores)

            # eventqueue
            event = FlowArrivalEvent(mcastFlow.startTime, flow)
            eventQueue.push(event, mcastFlow.startTime)

    def UpdateFlow(self, curFlow, eventType, eventQueue, curTime):
        # type: (object, object, object, object) -> object
        """
        Fariness bandwidth allocation
        Find related flows and update their flowBw, transTime, endTime.
        """
    
        if eventType == FlowArrival:
            self.runningFlows.append(curFlow)
            pathInLink = curFlow.pathLinkIds
            for link_id in pathInLink:
                link = self.Links[link_id]
                if curFlow.flowId not in link.flowIds:
                    link.flowIds.append(curFlow.flowId)
        
        elif eventType == FlowFinished:
            # update multicast flows
            mflowId = curFlow.mcastFlowId
            print curFlow
            if mflowId != 0:
                print 'multicast',mflowId
                mcastFlow = self.mcastFlows[mflowId]
                mcastFlow.dstHosts[curFlow.endId] = True
                
                if all(mcastFlow.dstHosts.values()):
                    mcastFlow.finishTime = curTime
                    print 'mcastFlow transmission time', mcastFlow.finishTime-mcastFlow.startTime
                else:
                    srcHost = curFlow.startId
                    dstHost = curFlow.endId
                    srcSubHosts = self.mcastSubtree[mflowId][srcHost]
                    dstSubHosts = self.mcastSubtree[mflowId][dstHost]
                    print srcSubHosts
                    print dstSubHosts
                    
                    if srcSubHosts:
                        index1 = len(srcSubHosts)/2
                        host1 = srcSubHosts[index1]
                        self.mcastSubtree[mflowId][srcHost] = srcSubHosts[:index1]
                        self.mcastSubtree[mflowId][host1] = srcSubHosts[index1+1:]
                        print 'host1', host1

                        flow = assignNewFlow(srcHost, host1, self.flows, mcastFlow.flowSize, curTime, \
                                            mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores)
                        self.runningFlows.append(flow)
                        pathInLink = flow.pathLinkIds
                        for link_id in pathInLink:
                            link = self.Links[link_id]
                            if flow.flowId not in link.flowIds:
                                link.flowIds.append(flow.flowId)
                    
                    if dstSubHosts:  
                        index2 = len(dstSubHosts)/2
                        host2 = dstSubHosts[index2]
                        self.mcastSubtree[mflowId][dstHost] = dstSubHosts[:index2]
                        self.mcastSubtree[mflowId][host2] = dstSubHosts[index2+1:]
                        flow = assignNewFlow(dstHost, host2, self.flows, mcastFlow.flowSize, curTime, \
                                            mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores)
                        self.runningFlows.append(flow)
                        pathInLink = flow.pathLinkIds
                        for link_id in pathInLink:
                            link = self.Links[link_id]
                            if flow.flowId not in link.flowIds:
                                link.flowIds.append(flow.flowId)
                        
                        print 'host2', host2

                    print self.mcastSubtree

            #
            curFlow.finishTime = curTime
            pathInLink = curFlow.pathLinkIds
            for link_id in pathInLink:
                link = self.Links[link_id]
                link.flowIds.remove(curFlow.flowId)
            
            self.runningFlows.remove(curFlow)
            self.finishedFlows.append(curFlow)
            print len(self.finishedFlows)
        
        elif eventType == Reconfigure:
            for flow in self.runningFlows:
                pathInLink = flow.pathLinkIds
                for linkId in pathInLink:
                    link = self.Links[linkId]
                    if flow.flowId not in link.flowIds:
                        link.flowIds.append(flow.flowId)
        
        current_flows = list(self.runningFlows)
        current_links = list(self.Links)
        residual_cap = {}
        residual_flowIds = {}
        for link_id in self.Links:
            link = self.Links[link_id]
            residual_cap[link_id] = link.linkCap
            residual_flowIds[link_id] = list(link.flowIds)
        
        self.runningFlowsDict = {}
        for flow in current_flows:
            self.runningFlowsDict[flow.flowId] = flow
            flow.remainSize -= flow.bw * (curTime - flow.updateTime)
            flow.remainTime -= curTime - flow.updateTime
            flow.bw = 0.0

            # update flow and eventQueue
            if flow.finishEvent != 0:
                flow.finishEvent.cancelled = True

        counter = 0
        while current_flows: # multi iterations
            bottleneck_rate = link.linkCap            
            for link_id in current_links:
                if len(residual_flowIds[link_id]) != 0:
                    if bottleneck_rate >= residual_cap[link_id] / len(residual_flowIds[link_id]):
                        bottleneck_rate = residual_cap[link_id] / len(residual_flowIds[link_id])
                else:
                    current_links.remove(link_id)
            bottleneck_rate = max([min([bottleneck_rate, link.linkCap]), 0.001])

            # update flow bandwidth
            for flow in current_flows:
                flow.bw += bottleneck_rate
            
            # delete bottleneck links and update related flows
            for link_id in current_links:
                residual_cap[link_id] -= bottleneck_rate * len(residual_flowIds[link_id])
                if residual_cap[link_id] <= 1.0e-20:
                    flowIds = list(residual_flowIds[link_id])
                    
                    for flow_id in flowIds:
                        flow = self.runningFlowsDict[flow_id]
                        pathInLink = flow.pathLinkIds
                        for flow_link_id in pathInLink:
                            residual_flowIds[flow_link_id].remove(flow_id)
                        
                        flow.updateTime = curTime
                        flow.finishTime = curTime + flow.remainSize/flow.bw + PROPAGATION_DELAY*len(flow.pathLinkIds)
                        event = FlowFinishedEvent(flow.finishTime, flow)
                        flow.finishEvent = event
                        eventQueue.push(event, event.time)

                        current_flows.remove(flow)
                    current_links.remove(link_id)
        
        del current_flows
        del current_links
        del residual_cap
        del residual_flowIds